#include "File/BOSSOutputLoader.h"
#include "Functions/Error.h"
#include "Functions/String.h"
#include "Functions/TerminalIO.h"
#include "Globals/Settings.h"
#include "TFile.h"
#include "TKey.h"
#include "TList.h"
#include "TObject.h"
#include "TSystemDirectory.h"
#include "TTree.h"
#include <float.h>
#include <iomanip>
#include <iostream>
#include <map>
#include <sstream>
#include <utility>
using namespace std;
using namespace BOSS_Afterburner::TerminalIO;
using namespace BOSS_Afterburner::String;

BOSSOutputLoader::BOSSOutputLoader(const YAML::Node& input)
{
  const auto& filename = input["Input file or directory"];
  if(!filename || !filename.IsScalar()) EXCEPTION("Missing key \"Input file or directory\"");
  Import(filename.as<string>().c_str());
}

BOSSOutputLoader::BOSSOutputLoader(const TString& path)
{
  Import(path);
}

void BOSSOutputLoader::Import(const TString& path)
{
  fFileNames.AddFilesRecursive(path);
  if(fFileNames.IsEmpty()) EXCEPTION(Form("Unable to load files from path \"%s\"", path.Data()));
  LoadChains();
}

/// Get the number of events in one of the `TChain`s.
Long64_t BOSSOutputLoader::GetEntries(const string& chainName)
{
  GetChain(chainName).GetEntries();
}

/// Return the largest number of events in all of the `TChain`s in the loaded ROOT file.
Long64_t BOSSOutputLoader::GetLargestEntries() const
{
  Long64_t largest = 0;
  for(auto it = fChains.begin(); it != fChains.end(); ++it)
    if(it->second.GetEntries() > largest) largest = it->second.GetEntries();
  return largest;
}

/// Check if file has chain with name `chainName`.
Bool_t BOSSOutputLoader::HasChain(const string& chainName) const
{
  if(fChains.find(chainName) == fChains.end()) return false;
  return true;
}

/// Print information about all trees in the `TFile`.
void BOSSOutputLoader::Print() const
{
  PrintFiles();
  PrintChains();

  // * Print contents of one input file * //
  TFile file(fFileNames.First().Data());
  if(file.IsZombie()) return;
  std::cout << endl;
  std::cout << "Output of TFile::ls():" << endl;
  std::cout << "---------------------" << endl;
  file.ls();
  std::cout << endl;
}

void BOSSOutputLoader::PrintFiles() const
{
  if(fFileNames.GetNFiles() > 1)
  {
    std::cout << "Loaded " << fFileNames.GetNFiles() << " files:" << endl;
    for(const auto& file : fFileNames.GetFiles())
      std::cout << "  \"" << file << "\"" << endl;
  }
  else if(fFileNames.GetNFiles())
  {
    std::cout << "Loaded file \"" << fFileNames.First() << "\"" << endl;
  }
  else
    std::cout << "No files loaded" << endl;
}

void BOSSOutputLoader::PrintChains() const
{
  if(fChains.size())
  {
    std::cout << "Contains " << fChains.size() << " TChains:" << endl;
    for(const auto& chain : fChains)
    {
      std::cout << "    \"" << chain.second.GetName() << "\"\t\"" << chain.second.GetTitle()
                << "\"\t" << chain.second.GetEntries() << " entries" << endl;
    }
  }
  else
    std::cout << "No TChains loaded" << endl;
}

void BOSSOutputLoader::PrintStructure()
{
  if(fChains.size())
  {
    std::cout << "Contains " << fChains.size() << " TChains:" << endl;
    for(auto& chain : fChains)
    {
      TString title = chain.second.GetTitle();
      std::cout << "  Name: \"" << chain.second.GetName() << "\" (" << chain.second.GetEntries()
                << ")" << endl;
      if(!title.EqualTo("ks N-Tuple example")) std::cout << "  Title: " << title << endl;
      auto arr = chain.second.GetListOfBranches();
      if(!arr) continue;
      for(auto i = 0; i < arr->GetEntries(); ++i)
        cout << "    " << arr->At(i)->GetTitle() << endl;
    }
  }
  else
    std::cout << "No TChains loaded" << endl;
}

/// Print information about a certain tree.
/// This function loops over the `unordered_map` of file names and over thes `unordered_map` of trees and prints its name and number of events. For each tree, information about its branches is also printed.
void BOSSOutputLoader::Print(const char* nameOfTree, Option_t* option)
{
  TChain* chain = &GetChain(nameOfTree);
  chain->Print(option);
  std::cout << "------------------------------------" << endl;
  std::cout << "Total number of events in TTree \"" << chain->GetName() << "\": " << scientific
            << chain->GetEntries() << endl;
  std::cout << endl;
}

/// Method that prints the number of entries of each `TChain` in the ROOT file. This allows you to get an idea of the cut flow of the original macro. (Of course, have a look at your original source code `.cxx` to see what these numbers mean.)
void BOSSOutputLoader::PrintCutFlow()
{
  if(!fChains.size()) return;
  /// -# Write number of entries per `TChain`.
  auto width = ceil(log10(GetLargestEntries()));
  width += 2;
  width += width / 3;
  auto orderedmap = CreateOrderedMap();
  std::cout << endl << "CUT FLOW FOR \"" << fFileNames.GetCommonPWD() << "\"" << endl;
  for(auto it : orderedmap)
  {
    TString chainName = it->GetName();
    if(chainName.EqualTo("_cutvalues")) continue; /// Do not print `"_cutvalues"` branch
    std::cout << right << setw(width) << CommaFormattedString(it->GetEntries()) << "  " << left
              << "\"" << it->GetName() << "\"" << endl;
  }
  std::cout << endl;
  /// -# If there is a `TChain` called `"_cutvalues"`, print these values and names as well using `PrintCuts`.
  PrintCuts();
}

/// This function creates a `list` of `TableEntry`s from the `"_cutvalues"` branch which is then printed. The function is backward compatible with older versions of BOSS IniSelect output.
Table BOSSOutputLoader::PrintCuts(bool print)
{
  /// <ol>
  /// <li> **Abort** if there is no `TChain` `"_cutvalues"` or if it is empty.
  auto key = fChains.find("_cutvalues");
  if(key == fChains.end()) return Table();
  if(!key->second.GetEntries()) return Table();
  auto& chain = key->second;

  /// <li> Determine version.
  enum Version
  {
    Error,
    Indexed,
    CountsOnly,
    MulptleOfThree
  };
  Version version = Error;
  if(MapHasKey(chain.Get<int>(), "index"))
    version = Indexed;
  else
  {
    if(true)
      // if(chain.GetEntries() == fFileNames.GetNFiles())
      version = CountsOnly;
    else if(chain.GetEntries() % 3 == 0)
      version = MulptleOfThree;
  }
  if(version == Error)
  {
    PrintWarning(
      Form("Unable to establish cut value version:\n  number of entries: %u\n  number of files: %u",
           chain.GetEntries(), fFileNames.GetNFiles()));
    return Table();
  }

  /// <li> Fill table.
  Table    table;
  TIter    next(chain.GetListOfBranches());
  TObject* obj = nullptr;
  while(obj = next())
  {
    const string branchName = obj->GetName();
    if(!branchName.compare("index")) continue;

    Long64_t count = 0;
    table << branchName;
    switch(version)
    {
      case Indexed:
      {
        /// <li> **CASE1**. If `"_cutvalues"` contains a branch called `"index"`, assume that the `TTree` contains double arrays.
        /// <ol>
        /// <li> `SetBranchAddress` for the array of this branch. **Abort** if `SetBranchAddress` does not return `0`.
        double values[3];
        auto   res = chain.SetBranchAddress(branchName.c_str(), &values);
        if(res)
        {
          std::cout << "ERROR: failed to load \"" << branchName << "\" in branch \"_cutvalues\""
                    << endl;
          std::cout << "  -->> SetBranchAddress returned " << res << endl;
          return Table();
        }
        /// <li> Get the total count value by adding up each 3rd entry of the array.
        for(int i = 0; i < chain.GetEntries(); ++i)
        {
          chain.GetEntry(i);
          count += values[2]; // total per file
        }
        /// <li> Entry `0`: this is information written from the `CutObject` to the first file of the `TChain` loaded here. We use it to load the `min` and `max` value (array entries 0 and 1).
        /// @see TrackSelector::AddAndWriteCuts
        chain.GetEntry(0);
        if(values[0] < -0.9 * DBL_MAX)
          table << string();
        else
          table << values[0]; // min
        if(values[1] > 0.9 * DBL_MAX)
          table << string();
        else
          table << values[1]; // min
        /// <li> Drop address allocated for this branch.
        chain.ResetBranchAddress(chain.GetBranch(branchName.c_str()));
        /// </ol>
        break;
      }
      case CountsOnly:
        /// <li> **CASE2**. If `"_cutvalues"` contains the same number of entries as the number of files in the `TChain`, **only** print a list of cut paramters and their values:
        for(int i = 0; i < chain.GetEntries(); ++i)
        {
          chain.GetEntry(i);
          try
          {
            count += chain.Get<double>(branchName);
          }
          catch(const out_of_range& e1)
          {
            try
            {
              count += chain.Get<int>(branchName);
            }
            catch(const out_of_range& e2)
            {
              count += chain.Get<unsigned long long>(branchName);
            }
          }
        }
        break;
      case MulptleOfThree:
      {
        /// <li> **CASE3**. If `"_cutvalues"` contains at least 3 entries, consider entry 0 to be the `min` value, entry 1 to be the `max`, and entry 2 to be `count` (the number of events or tracks that passed the cut).
        chain.GetEntry(0); // min
        table << chain.Get<double>(branchName);
        chain.GetEntry(1); // max
        table << chain.Get<double>(branchName);
        count = 0;
        for(Long64_t i = 2; i < chain.GetEntries(); i += 3)
        {
          chain.GetEntry(i); // count
          count += chain.Get<double>(branchName);
        }
        break;
      }
    }
    table << CommaFormattedString(count);
    table.SetLastSortValue(count);
    table.SetLastName(branchName);
    table << TableRow::endrow;
  }

  /// <li> Set headers.
  switch(version)
  {
    case Indexed:
    case MulptleOfThree:
      table.SetHeader({"CUT NAME", "MIN", "MAX", "TOTAL"});
      table.SetAlign(1, TableColumn::Internal);
      table.SetAlign(2, TableColumn::Internal);
      table.SetAlign(3, TableColumn::Right);
      break;
    case CountsOnly:
      table.SetHeader({"CUT NAME", "TOTAL"});
      table.SetAlign(1, TableColumn::Right);
      break;
  }

  /// <li> Print loaded values as a table.
  table.Sort();
  if(print) table.Print();

  /// <li> Reset branches to avoid a crash at the end of your executable.
  chain.ResetBranchAddresses();

  /// </ol>
  return table;
}

/// Draw the distributions of all branches available in the ROOT file.
TH1F* BOSSOutputLoader::Draw(const string& chainName, const char* branchX, const Int_t nBinx,
                             const double x1, const double x2, Option_t* option,
                             const TString& logScale, const char* cut)
{
  return GetChain(chainName).DrawHistogram(branchX, nBinx, x1, x2, option, true, logScale, cut);
}

/// Draw the distributions of all branches available in the ROOT file.
TH2F* BOSSOutputLoader::Draw(const string& chainName, const char* branchX, const char* branchY,
                             const Int_t nBinx, const double x1, const double x2, const Int_t nBiny,
                             const double y1, const double y2, Option_t* option,
                             const TString& logScale, const char* cut)
{
  return GetChain(chainName).DrawHistogram(branchX, branchY, nBinx, x1, x2, nBiny, y1, y2, option,
                                           true, logScale, cut);
}

TH1* BOSSOutputLoader::Draw(const string& chainName, const char* varexp, const char* cut,
                            Option_t* option, const TString& logScale)
{
  return GetChain(chainName).DrawHistogram(varexp, cut, option, true, logScale);
}

TH1* BOSSOutputLoader::Draw(const BranchPlotOptions& options)
{
  if(!options.IsOK()) return nullptr;
  if(!HasChain(options.TreeName()))
  {
    PrintWarning(Form("Cannot draw non-existing branch \"%s\" in file \"%s\"", options.TreeName(),
                      fFileNames.GetCommonPWD().Data()));
    return nullptr;
  }
  return GetChain(options.TreeName()).DrawHistogram(options);
}

TH1* BOSSOutputLoader::Draw(const YAML::Node& input)
{
  if(!input) return nullptr;
  const auto& treeDef = input["tree name"];
  if(!treeDef || !treeDef.IsScalar()) EXCEPTION("Missin key \"tree name\"");
  string tree = treeDef.as<string>();
  if(!HasChain(tree))
  {
    PrintWarning(Form("Cannot draw non-existing branch \"%s\" in file \"%s\"", tree.c_str(),
                      fFileNames.GetCommonPWD().Data()));
    return nullptr;
  }
  return GetChain(tree).DrawHistogram(input);
}

TH1* BOSSOutputLoader::Draw(const string& tree, const YAML::Node& input)
{
  if(!input) return nullptr;
  if(!HasChain(tree))
  {
    PrintWarning(Form("Cannot draw non-existing branch \"%s\" in file \"%s\"", tree.c_str(),
                      fFileNames.GetCommonPWD().Data()));
    return nullptr;
  }
  return GetChain(tree).DrawHistogram(input);
}

/// Draw the distributions of all branches available in the ROOT file.
void BOSSOutputLoader::DrawAndSaveAllBranches(const string& chainName, Option_t* option,
                                              const TString& logScale)
{
  GetChain(chainName).DrawAndSaveAllBranches(option, logScale);
}

/// Draw the distributions of all branches available in the ROOT file.
void BOSSOutputLoader::QuickDrawAndSaveAll(Option_t* option)
{
  for(auto it = fChains.begin(); it != fChains.end(); ++it)
    it->second.DrawAndSaveAllBranches(option);
}

/// Create a `list` that is ordered by the number of entries in the `TChain`.
list<ChainLoader*> BOSSOutputLoader::CreateOrderedMap()
{
  // * Create list of pointers to `ChainLoader` * //
  list<ChainLoader*> outputList;
  for(auto it = fChains.begin(); it != fChains.end(); ++it)
  {
    TString name(it->second.GetName());
    // if(!name.EqualTo("_cutvalues"))
    outputList.push_back(&(it->second));
  }
  // * Sort resulting list based on number of entries * //
  outputList.sort([](ChainLoader* const& a, ChainLoader* const& b) {
    string nameA(a->GetName());
    string nameB(b->GetName());
    return a->GetEntries() != b->GetEntries() ? a->GetEntries() > b->GetEntries() : nameA < nameB;
  });
  return outputList;
}

/// Create a list of `TChain`s based on the files in the directory you are loading.
/// If the ROOT `TFile` could be successfully loaded, this method will generate a `unordered_map` of pointers to all `TChain`s in the main folder of the ROOT file. The first file in the directory is used to make an inventory of all `TTree`s within the file.
void BOSSOutputLoader::LoadChains()
{
  /// -# Load first ROOT file `fFileNames` list.
  TFile    file(fFileNames.First().Data());
  TIter    next(file.GetListOfKeys());
  TObject* obj  = nullptr;
  TKey*    key  = nullptr;
  TTree*   tree = nullptr;
  while((obj = next()))
  {
    if(!(key = dynamic_cast<TKey*>(obj))) continue;
    if(!(tree = dynamic_cast<TTree*>(key->ReadObj()))) continue;
    fChains.emplace(tree->GetName(), tree);
    auto chain = &GetChain(tree->GetName());
    for(const auto& filename : fFileNames.GetFiles())
      chain->Add(filename.Data());
    chain->BookAddresses();
  }
  /// -# Print terminal output.
  std::cout << "Created a map of " << fChains.size() << " TChains from ";
  if(fFileNames.GetNFiles() > 1) std::cout << fFileNames.GetNFiles() << " files in common ";
  std::cout << "input path \"";
  if(fFileNames.GetNFiles() > 1)
    std::cout << fFileNames.GetCommonPWD();
  else
    std::cout << fFileNames.First();
  std::cout << "\"" << endl;
}